# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: jingyuan

import subprocess
import tempfile
import os
import json
import traceback
import warnings

import pandas as pd
from absl import app, flags
from coin.exchange.bybit.kr_rest.native_private_client import \
    BybitNativePrivateClient
from cc.appcoin2.strategy.support.constants import LEVERGAE_MAP_LIST
from cc.appcoin2.strategy.support.util import get_cmd_vault_key_copy


class BybitClient(BybitNativePrivateClient):
  def __init__(self, *args, **kwargs):
    self.settle_coin = kwargs.pop('settle_coin', 'USDT')
    super(BybitClient, self).__init__(*args, **kwargs)
    self.bins = {}

  def set_risklimit(self, symbols, limit):
    pos = self.positions(symbols)
    for symbol in symbols:
      risklimit = self.risk_bins([symbol]).sort_values('limit')
      risk_idx = risklimit[risklimit.limit > limit].index[0]
      risk_id = risklimit.loc[risk_idx].riskId
      this_limit = risklimit.loc[risk_idx].limit
      if symbol in pos.index and risk_id == pos.loc[symbol].riskId:
        print(f"{symbol} already at risk limit {this_limit}")
        continue
      resp = self._query('POST', '/contract/v3/private/position/set-risk-limit', {
          'symbol': symbol,
          'riskId': risk_id,
          'positionIdx': 0
      }).msg
      print(f"{symbol}(riskId: {risk_id}) -> max pos set to {this_limit}", resp)

  def positions(self, symbols):
    resp = self._query('GET', '/contract/v3/private/position/list', {
        'settleCoin': self.settle_coin
    }).msg
    pos = pd.DataFrame(resp['result']['list'])
    symbols = set(symbols).union(pos[pos['size'] != 0].symbol.unique())
    bins = self.risk_bins(symbols)
    columns = [
        'positionIdx', 'riskId', 'symbol', 'size', 'positionValue', 'entryPrice', 'leverage',
        'cumRealisedPnl', 'unrealisedPnl'
    ]
    df = pos[columns].merge(bins[['limit', 'maxLeverage', 'riskId', 'symbol']],
                            on=['riskId', 'symbol'])
    df = df[df.symbol.isin(symbols)].set_index(['symbol'])
    cols = [
        "positionIdx",
        "riskId",
        "limit",
        "leverage",
        "positionValue",
        "size",
        "entryPrice",
        "cumRealisedPnl",
        "unrealisedPnl",
        "maxLeverage",
    ]
    return df[cols]

  def risk_bins(self, symbols):
    dfs = []
    for symbol in symbols:
      if symbol not in self.bins:
        resp = self._query("GET", '/derivatives/v3/public/risk-limit/list', {'symbol': symbol})
        df = pd.DataFrame(resp.msg['result']['list']).rename({
            'id': "riskId"
        }, axis=1).astype({'limit': 'float'})
        self.bins[symbol] = df
      else:
        df = self.bins[symbol]
      dfs.append(df)
    return pd.concat(dfs)


def get_bybit_account(strategy):
  if os.path.exists(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'):
    driver = json.load(open(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'))
    for exch, config in driver['order'].items():
      if config['mea'] == 'Futures.Bybit.v3-linear':
        return config['connection_config']['key_filepath'].split('/')[3]
  raise ValueError(f"look up account for {strategy} failed")


def main(argv):
  pd.set_option('display.max_rows', None)
  if flags.FLAGS.strategy and not flags.FLAGS.account:
    account = get_bybit_account(flags.FLAGS.strategy)
    print(f'{flags.FLAGS.strategy} is using account {account}')
  else:
    account = flags.FLAGS.account
  symbols = [
      s + (flags.FLAGS.settle_coin if "USD" not in s else "")
      for s in flags.FLAGS.symbols.split(',')
  ]
  updated_symbols = []
  for s in symbols:
    for key in LEVERGAE_MAP_LIST:
      s = s.replace(key, LEVERGAE_MAP_LIST[key])
    updated_symbols.append(s)
  symbols = updated_symbols
  with tempfile.NamedTemporaryFile() as file:
    if flags.FLAGS.key_filepath:
      key_file = flags.FLAGS.key_filepath
    else:
      key_file = file.name
      vault_key_name = f"trade/{account}/bybit/trade_key.json"
      cmd = get_cmd_vault_key_copy(vault_key_name, key_file)
      output = subprocess.check_output(cmd, shell=True)
    client = BybitClient(key_file=key_file, settle_coin=flags.FLAGS.settle_coin)

    results = client.positions(symbols)
    print("\n#####################")
    print("-- Before Adjust --")
    print(results)
    print("######################\n")


    if not flags.FLAGS.query_only:
      try:
        client.set_risklimit(updated_symbols, flags.FLAGS.maxpos)
      except BaseException:
        traceback.print_exc()
        warnings.warn("Failed to set risklimit.")
      print("\n#####################")
      print("-- After Adjust --")
      print(client.positions(updated_symbols))
      print("######################\n")


if __name__ == '__main__':
  flags.DEFINE_string('strategy', '', 'strategy name')
  flags.DEFINE_string('account', '', 'account name')
  flags.DEFINE_string('symbols', 'ZILUSDT', '')
  flags.DEFINE_string('settle_coin', 'USDT', 'USDT or USDC')
  flags.DEFINE_string('key_filepath', '', 'key filepath')
  flags.DEFINE_integer('maxpos', 0, 'required max exposure in USD')
  flags.DEFINE_bool('query_only', True, 'check only')
  app.run(main)

# ./pyrunner cc/appcoin2/strategy/support/bybit_leverage.py  --account=xunke02 --query_only=0 --symbols=LTCUSDT --maxpos=500000
